from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 初始化DeepSeek-R1模型与Tokenizer（容器内部署）
model_name = "deepseek-ai/deepseek-7b-r1"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype=torch.float16,  # 半精度推理，减少GPU占用
    device_map="auto"  # 自动分配设备（GPU优先）
)
model.eval()  # 推理模式，禁用Dropout


def generate_answer(query: str, retrieved_docs: list):
    # 构建Prompt：结合用户请求与检索结果
    prompt = f"""基于以下参考信息，回答用户的问题：
    参考信息：
    {chr(10).join([f"- {doc}" for doc in retrieved_docs])}
    用户问题：{query}
    要求：1. 仅基于参考信息回答，不添加外部知识；2. 语言简洁，逻辑清晰；3. 若参考信息无相关内容，回复“暂无相关信息”。
    """

    # 模型推理
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():  # 禁用梯度计算，提升速度
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,  # 最大生成长度
            temperature=0.7,  # 随机性控制（0.7适合问答）
            top_p=0.9,
            do_sample=True,
            eos_token_id=tokenizer.eos_token_id
        )

    # 解析生成结果
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # 提取回答部分（去除Prompt）
    answer = answer.split("用户问题：")[1].split("要求：")[0].strip()
    return answer


# 示例：结合检索结果生成回答
if __name__ == "__main__":
    query = "产品A的定价是多少"
    retrieved_docs = ["产品A的定价为999元，支持30天无理由退货", "产品A的核心功能包括智能识别、自动同步、多设备兼容"]
    answer = generate_answer(query, retrieved_docs)
    print("生成回答：", answer)
    # 输出：生成回答：产品A的定价为999元，支持30天无理由退货。